clear all
colors={[0 0.4470 0.7410],[0.8500 0.3250 0.0980],[0.9290 0.6940 0.1250],[0.4940 0.1840 0.5560],[0.4660 0.6740 0.1880],[0.3010 0.7450 0.9330],[0.6350 0.0780 0.1840]};


pp = load('mesh_delay_scan_R0_3e_5MHz.mat');
Omega_m = pp.Omega_m;
Delta_m = pp.Delta_m;
mu_m = pp.mu_m;
pg_m = pp.pg_m;
count_m = pp.count_m;

P = [2 1 3];
X = permute(Omega_m,P);
Y = permute(Delta_m,P);
Z = permute(mu_m,P);
Vg = permute(pg_m,P);
Vc = permute(count_m,P);
global F_pg F_count
F_pg = griddedInterpolant(X,Y,Z,Vg,'makima');
F_count = griddedInterpolant(X,Y,Z,Vc,'makima');


% unit of time is us
% unit of energy is MHz 
global v0 v1 v2
v0 = 4.18; %m/s
v1 = 3.45; %m/s
v2 = 0.74; %m/s

% % initial molecule state is n=1
v0 = 4.81; %m/s
v1 = 4.19; %m/s
v2 = 2.49; %m/s


%% data



% R0 '20230513_100VR_30G_10kHz_REMPI_scan_4'

x0 = [-220  -160  -150  -140   -90   -40     0    10    15    20    30    40    50    90];
signal = [1.9500   10.3684   14.4211   18.2632   16.7500   18.8000   18.4211   19.5500   18.2632   18.9474   21.0500   12.8947    3.0000    1.9500];
e_signal = [ 0.4213    0.7572    0.9146    1.0083    0.9631    0.9975    1.0421    1.0404    1.0246    1.0367    1.0665    0.8823    0.5130    0.3708];

bk = [1.7500    3.5263    4.7895    4.2105    3.7500    4.9000    4.5789    4.4500    5.0000    4.7368    3.8000    3.4737    1.3158    2.3500];
e_bk = [0.4093    0.4618    0.5741    0.5263    0.5268    0.5477    0.5978    0.5723    0.5931    0.5717    0.5244    0.5316    0.4178    0.3969];

% % R0 '20230622_100VR_30G_10kHz_SR_scan_5'
% 
% x0=[ -220  -160  -150  -140   -90   -40     0    10    15    20    30    40    50    90];
% signal = [ 1.0385    4.1154    9.6154   10.1923    8.6667    9.6154   10.3462   10.7692   10.9231    9.6923    9.5000    4.7692    1.2308    0.3846];
% e_signal = [0.3241    0.5059    0.6504    0.6695    0.6264    0.6772    0.6869    0.6815    0.6944    0.6572    0.6538    0.4679    0.3649    0.2433];
% 
% bk = [1.0769    1.4231    2.2692    3.1538    2.7778    2.0769    2.4231    3.1923    2.6538    2.9615    2.6538    1.7692    0.3846    0.9615];
% e_bk = [0.3264    0.3903    0.3749    0.4213    0.4174    0.4107    0.4089    0.4160    0.4052    0.4160    0.4052    0.3218    0.3172    0.2852];

% % 2G R0 '20230626_100VR_2G_10kHz_SR_scan_3'
% x0 = [-220  -160  -150  -140   -90   -40     0    10    15    20    30    40    50    90];
% signal = [0.8235    6.7647    9.0000   10.2941   10.2941   11.0000   11.2353   10.5294    9.6471    9.5882    9.4706    6.0000    1.7059    0.9412];
% e_signal = [0.4242    0.6733    0.7737    0.8422    0.8172    0.8625    0.8545    0.8463    0.8151    0.8340    0.8001    0.6903    0.4362    0.3328];
% bk = [ 0.4118    2.1765    2.2941    1.8235    2.1765    2.1176    3.2941    2.5882    3.2353    2.3529    2.4706    1.5294    1.4706    1.9412];
% e_bk = [0.3946    0.4282    0.4518    0.4594    0.4362    0.4706    0.5128    0.4991    0.5359    0.5195    0.4779    0.4242    0.4201    0.4118];
% 
y0 = signal - bk;
ey0 = sqrt(e_signal.^2+e_bk.^2);

% % 30G 22 + N=1,
% x0 = [-220  -160  -150  -140   -90   -40     0    10    15    20    30    40    50    90];
% y0 = [1.2000    3.2000    4.4000    5.0000    6.0000    5.6000    6.4000    3.2000    8.0000    4.2000    5.2000    2.6000    1.2000    0.6000];
% ey0 = [0.7483    1.0583    1.0583    1.1136    1.1662    1.1662    1.2000    0.8944    1.2649    0.9592    1.0954    0.8246    0.5657    0.5292];
% y0 = y0;

% % 30G, 22 + N=1,
% % '20230628_100VR_30G_10kHz_SR_scan_34','20230628_100VR_30G_10kHz_SR_scan_12'
% x0 = [-220  -160  -150  -140   -90   -40     0    10    15    20    30    40    50    90];
% y0 = [ 0.8286    3.1429    5.0000    5.8857    5.3429    6.8000    6.8000    6.3143    6.9714    5.6571    5.7429    3.4571    1.3714    1.0571];
% ey0 = [ 0.2268    0.3429    0.4010    0.4408    0.4209    0.4625    0.4625    0.4598    0.4695    0.4238    0.4417    0.3603    0.2356    0.2231];
% 
% y0 = y0-1;


x0 = (-x0+30+20+12-3)*1e-3;
[x0,I] = sort(x0);
y0 = y0(I);
ey0 = ey0(I);

x0= x0(2:end);
y0 = y0(2:end);
ey0 = ey0(2:end);

%
% plot(xft,MC_decay1([10,10,10],xft),'-o','linewidth',2,'markersize',8,'displayname','n=1')
% plot(xft,MC_decay0([10,10,10],xft),'-o','linewidth',2,'markersize',8,'displayname','n=0')






%% fit R0
if 1
figure(46)
clf
hold on

errorbar(x0,y0,ey0,'o','linewidth',2,'markersize',8,'displayname','n=0')

xft = x0;

% fit_val = MC_delay0([0.18 18 10],xft);
% plot(xft,fit_val,'-','linewidth',2,'displayname','fitting');
options = optimset('Display','iter','MaxIter',100)
fun = @(c) sum((MC_delay0(c,x0)-y0').^2);
c0 = [0.33,13.8];
c0 = [0.26 9.6];
% c0 = [0.18 ];
c = fminsearch(fun,c0,options)

fit_val = MC_delay0(c,xft);
plot(xft,fit_val,'-o','markersize',8,'linewidth',2,'displayname','fitting');
cerr = extract_error_bar(@(c,t) MC_delay0(c,t),c,fit_val,y0',x0)

legend
xlabel('532 center delay (us)')
ylabel('KRb ion counts')
set(gcf,'color','white')
set(gca,'fontsize',16)

end


% function y = MC_delay2(c,t)
% global v2
%     n0 = c(1);
%     Omega = c(2);
%     tau = c(3);
% %     y = n0.*counts_int_angle(t,0.74,Omega,5).*exp(-t.^0.5/tau);
%     y = n0.*counts_int_angle(t,v2,Omega,5,tau);
% end

% function y = MC_delay1(c,t)
% global v1
%     n0 = c(1);
%     Omega = 14.48;
%     tau = c(2);
% %     y = n0.*counts_int_angle(t,v1,Omega,5).*exp(-t.^0.5/tau);
%     y = n0.*counts_int_angle(t,v1,Omega,5,tau);
% end

function y = MC_delay0(c,t)
global v0
    n0 = c(1);
    Omega = c(2);%10.3;9*13.7/10; %c(2);
%     tau = c(3);
    delta = -10;
    
%     y = n0.*counts_int_angle(t,v0,Omega,5).*exp(-t./tau);
    y = n0.*counts_int_angle0(t,v0,Omega,delta);
end

% function ctot = counts_int_angle(pulse_duration,v,Rabi0,Delta0)  % measured counts at a certain pulse duration
% %     rng(1);
%     avg = 1000;
%     cor = normrnd(0,1,[avg,3]);
%     cor = cor./sqrt(sum(cor.^2,2));
%     theta = acos(cor(:,3)); 
%     phi = atan(cor(:,2)./cor(:,1));
% %     t0 = rand(1,1000)*100e-6;
% %     r0 = v*t0;
%     t_dark = abs(125e-6./(v.*sin(theta).*sin(phi))); %now in unit of s
%     
%     
%     r0 = v*t_dark;
%     ctot=0;
% %     theta = pi/2;
%     for i = 1:length(theta)
%         th= theta(i);
%         N = round(3.0e-3./sin(th)./(v*100e-6));
%         
%         ctot = ctot+counts(N,pulse_duration,v,th,Rabi0,Delta0,r0(i)*sin(th),t_dark(i));
%     end
% %     ctot = ctot / avg;
% end

function ctot = counts_int_angle0(pulse_duration,v,Rabi0,Delta0)  % measured counts at a certain pulse duration
%     rng(1);
    avg = 1000;
    cor = normrnd(0,1,[avg,3]);
    cor = cor./sqrt(sum(cor.^2,2));
    theta = acos(cor(:,3)); 
    phi = atan(cor(:,2)./cor(:,1));
    t0 = rand(1,avg)*100e-6;
%     r0 = v*t0;
    t_dark = abs(125e-6./(v.*sin(theta).*sin(phi))); %now in unit of s
    
    
    r0 = v*(t_dark+t0');
    ctot=0;
%     theta = pi/2;
    for i = 1:length(theta)
        th= theta(i);
        N = round(3.0e-3./sin(th)./(v*100e-6));
        
        ctot = ctot+counts0(N,pulse_duration,v,th,Rabi0,Delta0,r0(i)*sin(th),t_dark(i));
    end
%     ctot = ctot / avg;
end

function counts = counts(N,mu,v,theta,Rabi0,Delta0,r0)  % measured counts at a certain pulse duration
% N is number of pulses.
%     n0 = round(td./(100e-6));
%     p_g_0 = exp(-n0.*pulse_duration.^0.5/tau).*((100-pulse_duration)/100).^2;
    p_g_0 = 1;
    counts = 0;
    p_g = 1;
    rrabi = Rabi(0:N',v*sin(theta),r0);
    [p_e_meta,p_g1_meta] = populations(rrabi,mu,v,theta,Rabi0,Delta0);  
%     size(p_e_meta)
%     size(p_g1_meta)
    for i = 1:N+1
%         counts = counts + p_e_meta(:,i).*p_g;
        counts = counts + p_e_meta(:,i).*p_g.*rrabi(i).^2;
%         counts = counts + p_e_meta(i).*p_g;
        p_g = p_g1_meta(:,i).*p_g; % remaining populations.
    end
    counts = counts.*p_g_0;
    
end

function counts = counts0(N,mu,v,theta,Rabi0,Delta0,r0,td)  % measured counts at a certain pulse duration
% N is number of pulses.
    n0 = round(td./(100e-6));
%     p_g_0 = exp(-n0.*pulse_duration.^0.5/tau).*exp(-pulse_duration/tau2);
%     p_g_0 = exp(-1e12*Rabi0^2*2*pi*0.005*(n0*0.2*1e-6)/(15.6e6+4*50*50*1e6/15.6));
    p_g_0 = 1; %exp(-pulse_duration/tau2);
    counts = 0;
    p_g = 1;
    rrabi = Rabi(0:N',v*sin(theta),r0);
    [p_e_meta,p_g1_meta] = populations(rrabi,mu,v,theta,Rabi0,Delta0);  
%     size(p_e_meta)
%     size(p_g1_meta)
    for i = 1:N+1
        counts = counts + p_e_meta(:,i).*p_g.*rrabi(i);
%         counts = counts + p_e_meta(i).*p_g;
        p_g = p_g1_meta(:,i).*p_g;%.*exp(-1e12*Rabi0^2.*rrabi(i).^2*2*pi*0.0001*100e-6/(15.6e6+4*50*50*1e6/15.6)); % remaining populations.
    end
    counts = counts.*p_g_0;

end

function [p_ee,p_gg] = populations(Rabi,mu,v,theta,Rabi0,Delta0) 
    % t1 is the delay between 670 REMPI on edge, and t1 is the 670 pulse on time 
    % hard code t1 as 30ns, 

    Omega = Rabi0*Rabi;
    % detuning from Doppler shift
    Delta = v*cos(theta)/3e8*4.49844e5+Delta0;

    ssz = [length(Omega),length(mu)];
    p_ee = Ion_count(repmat(Omega',[1,ssz(2)]),repmat(Delta,ssz),repmat(mu,[ssz(1),1]));
    p_gg = Pg(repmat(Omega',[1,ssz(2)]),repmat(Delta,ssz),repmat(mu,[ssz(1),1]));

end
  
function R = Rabi(n,v,r0)
% n is number of pulses 
    w = 1e-3; % beam waist
    r = r0+n.*v*100e-6;
    I = exp(-2*r.^2./w.^2);
    R = exp(-1*r.^2./w.^2);
end


function [se] = extract_error_bar(func_hand,c_pred,y_pred,y_data,x_data)
    % degrees of freedom in the problem
    dof = length(y_pred) - length(c_pred);

    % standard deviation of the residuals
    sdr = sqrt(sum((y_pred - y_data).^2)/dof);

    % jacobian matrix
    [J,~] = jacobianest(func_hand,c_pred,x_data);

    Sigma = sdr^2*inv(J'*J);

    % Parameter standard errors
    se = sqrt(diag(Sigma))';

%     % which suggest rough confidence intervalues around
%     % the parameters might be...
%     abupper = abfinal + 2*se
%     ablower = abfinal - 2*se
end


function counts1 = Ion_count(Omega,Delta,mu)
    global F_count
    counts1 = F_count(Omega',Delta',mu');
end

function pg = Pg(Omega,Delta,mu)
    global F_pg
    pg = F_pg(Omega',Delta',mu');
% pg = zeros(size(Omega));    
% inds = (Omega*2*pi).^2./(15.6*2*pi+4*(Delta*2*pi).^2./(15.6*2*pi))<=9/15*2*pi;
%     pg(inds) = exp(-length(inds).*(Omega(inds)*2*pi).^2../(15.6*2*pi+4*(Delta(inds)*2*pi).^2./(15.6*2*pi)));
%     pg(~inds) = F_pg(Omega(~inds),Delta(~inds),length(~inds));    
end
